package com.vue.util;

import com.vue.entity.vo.SseEmitterUTF8;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * @author: wrj
 * @since: 2024/11/27 11:57
 * @description:
 */
@Slf4j
public class SseEmitterUtil {

    /**
     * 使用map对象，便于根据userId来获取对应的SseEmitter，或者放redis里面
     */
    private final static Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    public static SseEmitterUTF8 connect(String userId) throws IOException{
        // 设置超时时间，0表示不过期。默认30S，超时时间未完成会抛出异常：AsyncRequestTimeoutException
        SseEmitterUTF8 sseEmitter = new SseEmitterUTF8(0L);
        // 注册回调
        sseEmitter.send(userId+"连接成功");
        sseEmitter.onCompletion(completionCallBack(userId));
        sseEmitter.onError(errorCallBack(userId));
        sseEmitter.onTimeout(timeoutCallBack(userId));
        sseEmitterMap.put(userId, sseEmitter);

        log.info("创建新的 SSE 连接，当前用户 {}, 连接总数 {}", userId, sseEmitterMap.size());
        return sseEmitter;
    }

    /**
     * 给制定用户发送消息
     *
     * @param userId 指定用户名
     * @param message 消息
     */
    public static void sendMessage(String userId, String message) {
        byte[] bytes = message.getBytes(StandardCharsets.UTF_8);

        if (sseEmitterMap.containsKey(userId)) {
            try {
                sseEmitterMap.get(userId).send(bytes, MediaType.TEXT_PLAIN);
                log.info("用户 {} 推送消息 {}", userId, message);
            } catch (IOException e) {
                log.error("用户 {} 推送消息异常", userId, e);
                removeUser(userId);
            }
        } else {
            log.error("消息推送 用户 {} 不存在，链接总数 {}", userId, sseEmitterMap.size());
        }
    }

    /**
     * 群发消息
     */
    public static void batchSendMessage(String message, List<String> ids) {
        ids.forEach(userId -> sendMessage(userId, message));
    }

    /**
     * 群发所有人
     */
    public static void batchSendMessage(String message) {
        sseEmitterMap.forEach((k, v) -> {
            try {
                v.send(message, MediaType.APPLICATION_JSON);
            } catch (IOException e) {
                log.error("用户 {} 推送异常", k, e);
                removeUser(k);
            }
        });
    }

    /**
     * 移除用户连接
     *
     * @param userId 用户 ID
     */
    public static void removeUser(String userId) {
        if (sseEmitterMap.containsKey(userId)) {
            sseEmitterMap.get(userId).complete();
            sseEmitterMap.remove(userId);
            log.info("移除用户 {}, 剩余连接 {}", userId, sseEmitterMap.size());
        } else {
            log.error("消息推送 用户 {} 已被移除，剩余连接 {}", userId, sseEmitterMap.size());
        }
    }

    /**
     * 获取当前连接信息
     *
     * @return 所有的连接用户
     */
    public static List<String> getIds() {
        return new ArrayList<>(sseEmitterMap.keySet());
    }

    /**
     * 获取当前的连接数量
     *
     * @return 当前的连接数量
     */
    public static int getUserCount() {
        return sseEmitterMap.size();
    }

    private static Runnable completionCallBack(String userId) {
        return () -> {
            log.info("用户 {} 结束连接", userId);
        };
    }

    private static Runnable timeoutCallBack(String userId) {
        return () -> {
            log.error("用户 {} 连接超时", userId);
            removeUser(userId);
        };
    }

    private static Consumer<Throwable> errorCallBack(String userId) {
        return throwable -> {
            log.error("用户 {} 连接异常", userId);
            removeUser(userId);
        };
    }
}
